from .base import FactorBase
from common.clickhouse_query import default_get_data
from common.variables import TimeFrameType
import numpy as np
import pandas as pd
from scipy.special import comb
import talib
from tqdm import tqdm
from typing import Literal


class RVI(FactorBase):
    def __init__(self, N: int = 14):
        self.N = N
        self.preloads = N

    @property
    def bfname(self):
        return f"{self.basename}_{self.N}"

    def calculate(
        self,
        symbol: str,
        timeframe: TimeFrameType,
        start: str,
        end: str,
        callbacks: list[callable],
    ) -> pd.DataFrame:
        df = default_get_data(
            symbol,
            timeframe,
            start,
            end,
            self.preloads,
            features=["open", "high", "low", "close"],
        )
        if len(df) < self.preloads:
            return None
        N = self.N
        a = df["close"] - df["open"]
        e = df["high"] - df["low"]
        NUMERATOR = (a + 2 * a.shift(1) + 2 * a.shift(2) + a.shift(3)) / 6
        DENOMINATOR = (e + 2 * e.shift(1) + 2 * e.shift(2) + e.shift(3)) / 6
        sma_NUMERATOR = talib.SMA(NUMERATOR, timeperiod=N)
        sma_DENOMINATOR = talib.SMA(DENOMINATOR, timeperiod=N)
        r = sma_NUMERATOR / sma_DENOMINATOR
        value = (r + 2 * r.shift(1) + 2 * r.shift(2) + r.shift(3)) / 6
        dfo = pd.DataFrame({"rvi": value}, index=df.index)
        return dfo


class Pascal(FactorBase):

    def __init__(self, N: int = 14):
        self.N = N
        self.preloads = N

    @property
    def bfname(self):
        return f"{self.basename}_{self.N}"

    def calculate(
        self,
        symbol: str,
        timeframe: TimeFrameType,
        start: str,
        end: str,
        callbacks: list[callable],
    ) -> pd.DataFrame:
        df = default_get_data(symbol, timeframe, start, end, self.preloads, ["close"])
        if len(df) < self.preloads:
            return None
        N = self.N
        weights = np.array([comb(N, k) for k in range(N)])
        weights = weights / weights.sum()
        result = np.convolve(df["close"], weights, mode="same")
        dfo = pd.DataFrame({"pascal": result}, index=df.index)
        return dfo


class CFO(FactorBase):
    def __init__(
        self,
        N: int = 14,
        kind: Literal["O", "H", "L", "C", "HL", "HLC", "HLCC", "OHLC"] = "C",
    ):
        self.N = N
        self.kind = kind
        self.preloads = N

    @property
    def bfname(self):
        return f"{self.basename}_{self.N}_{self.kind}"

    @staticmethod
    def get_cfo(prices: pd.Series):
        X = np.arange(len(prices))
        y = prices
        coes = np.polyfit(X, y, 1)
        forecast = np.polyval(coes, [len(prices)])
        return forecast

    def calculate(
        self,
        symbol: str,
        timeframe: TimeFrameType,
        start: str,
        end: str,
        callbacks: list[callable],
    ) -> pd.DataFrame:
        df = default_get_data(
            symbol,
            timeframe,
            start,
            end,
            self.preloads,
            ["open", "high", "low", "close"],
        )
        if len(df) < self.preloads:
            return None
        N = self.N
        kind = self.kind
        dc = {"O": df["open"], "C": df["close"], "H": df["high"], "L": df["low"]}
        value = dc[kind[0]]
        for x in kind[1:]:
            value = value + dc[x]
        value: pd.Series = value / len(kind)
        result = value.rolling(window=N, center=False).apply(self.get_cfo)
        value = (df["close"] - result) / df["close"]
        dfo = pd.DataFrame({"cfo": value}, index=df.index)
        return dfo


class TalibHT(FactorBase):
    def __init__(self, N: int = 14):
        self.N = N
        self.preloads = N

    @property
    def bfname(self):
        return f"{self.basename}_{self.N}"

    def calculate(
        self,
        symbol: str,
        timeframe: TimeFrameType,
        start: str,
        end: str,
        callbacks: list[callable],
    ) -> pd.DataFrame:
        df = default_get_data(
            symbol,
            timeframe,
            start,
            end,
            self.preloads,
            ["open", "high", "low", "close", "volume"],
        )
        if len(df) < self.preloads:
            return None
        self.N
        dfo = pd.DataFrame(
            {
                "HT_DCPERIOD": talib.HT_DCPERIOD(df["close"]),
                "HT_DCPHASE": talib.HT_DCPHASE(df["close"]),
                **{
                    x: y
                    for x, y in zip(
                        ["phasor_amplitude", "phasor_phase"],
                        talib.HT_PHASOR(df["close"]),
                    )
                },
                **{
                    x: y
                    for x, y in zip(
                        ["sine_wave", "lead_sine_wave "], talib.HT_SINE(df["close"])
                    )
                },
                "HT_TRENDMODE": talib.HT_TRENDMODE(df["close"]),
            },
            index=df.index,
        )
        return dfo


class TalibMomentum(FactorBase):
    def __init__(self, N: int = 14):
        self.N = N
        self.preloads = N

    @property
    def bfname(self):
        return f"{self.basename}_{self.N}"

    def calculate(
        self,
        symbol: str,
        timeframe: TimeFrameType,
        start: str,
        end: str,
        callbacks: list[callable],
    ) -> pd.DataFrame:
        df = default_get_data(
            symbol,
            timeframe,
            start,
            end,
            self.preloads,
            ["open", "high", "low", "close", "volume"],
        )
        if len(df) < self.preloads:
            return None
        self.N
        dfo = pd.DataFrame(
            {
                "HT_DCPERIOD": talib.HT_DCPERIOD(df["close"]),
                "HT_DCPHASE": talib.HT_DCPHASE(df["close"]),
                **{
                    x: y
                    for x, y in zip(
                        ["phasor_amplitude", "phasor_phase"],
                        talib.HT_PHASOR(df["close"]),
                    )
                },
                **{
                    x: y
                    for x, y in zip(
                        ["sine_wave", "lead_sine_wave "], talib.HT_SINE(df["close"])
                    )
                },
                "HT_TRENDMODE": talib.HT_TRENDMODE(df["close"]),
            },
            index=df.index,
        )
        return dfo


class TsFresh(FactorBase):
    def __init__(self, N: int = 14):
        self.N = N
        self.preloads = N

    @property
    def bfname(self):
        return f"{self.basename}_{self.N}"

    def calculate(
        self,
        symbol: str,
        timeframe: TimeFrameType,
        start: str,
        end: str,
        callbacks: list[callable],
    ) -> pd.DataFrame:
        df = default_get_data(
            symbol,
            timeframe,
            start,
            end,
            self.preloads,
            ["open", "high", "low", "close", "volume"],
        )
        if len(df) < self.preloads:
            return None
        import pandas as pd
        from tsfresh import extract_features
        from tsfresh.feature_extraction import EfficientFCParameters

        N = self.N
        results = []
        for x in tqdm(df.rolling(window=N, min_periods=N)):
            if len(x) < N:
                continue
            tmp = extract_features(
                x[["time", "close", "symbol"]],
                column_id="symbol",  # 组 ID
                column_sort="time",  # 时间排序
                default_fc_parameters=EfficientFCParameters(),  # 使用 Minimal 计算部分特征，可换成 EfficientFCParameters() 计算完整特征
            )
            results.append(tmp)
        features = pd.concat(results)
        features.index.name = "symbol"
        features.reset_index(inplace=True)
        return features
